import sys
import random
import json
import os

from rainforeLearn.gomoku.v2.agents.dqn_agent import GomokuDQNAgent
from rainforeLearn.gomoku.v2.configs.config import GomokuDQNConfig
from rainforeLearn.gomoku.v2.environment.env import GomokuEnvironment

# 尝试导入 Pygame，如果失败则打印提示
try:
    import pygame
except ImportError:
    print("Error: Pygame is not installed. Please install it using 'pip install pygame'")

    # 提供一个假的 pygame 模块，以便其他代码可以继续执行但会跳过图形部分
    class MockPygame:
        def init(self): pass

        def display(self):
            class MockDisplay:
                def set_mode(self, size): return None
                def set_caption(self, caption): pass
                def flip(self): pass
                def get_active(self): return None
            return MockDisplay()

        def event(self):
            class MockEvent:
                def get(self): return []
                QUIT = -1  # Dummy value
                MOUSEBUTTONDOWN = -2  # Dummy value
            return MockEvent()

        def time(self):
            class MockTime:
                def Clock(self):
                    class MockClock:
                        def tick(self, fps): pass
                    return MockClock()
            return MockTime()

        def font(self):
            class MockFont:
                def Font(self, name, size):
                    class MockRender:
                        def render(self, text, antialias, color): return MockRender()
                        def get_rect(self, center): return MockRender()
                        def blit(self, surface, rect): pass
                    return MockRender()
            return MockFont()

        def quit(self): pass
        K_ESCAPE = -3  # Dummy value

        def draw(self):
            class MockDraw:
                def line(self, s, c, p1, p2, w): pass
                def circle(self, s, c, p, r, w=0): pass
            return MockDraw()

    pygame = MockPygame()


class GomokuGameGUI:
    """五子棋GUI游戏"""

    def __init__(self, config: GomokuDQNConfig, model_path: str = None, ai_plays_first: bool = False):
        # 检查 Pygame 是否可用
        if 'pygame' not in sys.modules or pygame.display.get_active() is None:
            print("Pygame is not initialized or display not available. GUI will not start.")
            self.gui_enabled = False
            return

        self.gui_enabled = True
        try:
            print("🎮 Initializing Enhanced Gomoku Game...")
            pygame.init()
            print("✅ Pygame initialized.")
        except Exception as e:
            print(f"❌ Error initializing Pygame: {e}. GUI will not start.")
            self.gui_enabled = False
            return

        self.config = config
        self.board_size = self.config.board_size
        self.cell_size = 40
        self.padding = 20
        self.screen_width = self.board_size * self.cell_size + self.padding * 2
        self.screen_height = self.board_size * self.cell_size + self.padding * 2 + 80

        try:
            self.screen = pygame.display.set_mode((self.screen_width, self.screen_height))
            pygame.display.set_caption("Enhanced 五子棋 (Human vs Advanced AI)")
            print(f"🖥️  Screen created: {self.screen_width}x{self.screen_height}")
        except Exception as e:
            print(f"❌ Error creating screen: {e}")
            self.gui_enabled = False
            pygame.quit()
            return

        self.font = pygame.font.Font(None, 28)
        self.big_font = pygame.font.Font(None, 36)

        # 初始化游戏环境
        self.env = GomokuEnvironment(self.config)
        self.state = self.env.reset()

        # 玩家设置 - 修复玩家角色分配
        self.ai_plays_first = ai_plays_first
        if ai_plays_first:
            self.human_player = -1  # 人类玩白棋
            self.ai_player = 1      # AI玩黑棋
            self.current_turn_player = 1  # AI先手（黑棋）
            print("🤖 AI plays BLACK (goes first), Human plays WHITE")
        else:
            self.human_player = 1   # 人类玩黑棋
            self.ai_player = -1     # AI玩白棋
            self.current_turn_player = 1  # 人类先手（黑棋）
            print("👤 Human plays BLACK (goes first), AI plays WHITE")

        # 游戏状态
        self.game_over = False
        self.winner = None
        self.last_action = None

        # 统计信息
        self.game_stats = {
            'human_wins': 0,
            'ai_wins': 0,
            'draws': 0,
            'total_games': 0
        }

        # 加载统计信息
        self.load_stats()

        # 颜色定义
        self.BOARD_COLOR = (245, 222, 179)
        self.LINE_COLOR = (139, 69, 19)
        self.BLACK_STONE_COLOR = (20, 20, 20)
        self.WHITE_STONE_COLOR = (250, 250, 250)
        self.LAST_MOVE_COLOR = (255, 0, 0)
        self.LEGAL_MOVE_COLOR = (0, 255, 0)

        # AI设置
        self.ai_agent = None
        self.ai_enabled = False
        self.show_ai_thinking = True
        self.ai_thinking_time = 0.5

        if model_path:
            try:
                self.ai_agent = GomokuDQNAgent(self.config)
                self.ai_agent.load(model_path)
                self.ai_enabled = True
                print("✅ AI agent loaded successfully")
            except Exception as e:
                print(f"❌ Error initializing AI agent: {e}")
                self.ai_enabled = False
        else:
            print("📝 No AI model provided, human vs human mode")

        print("🚀 Game ready to start!")

    def load_stats(self):
        """加载游戏统计"""
        try:
            if os.path.exists('game_stats.json'):
                with open('game_stats.json', 'r') as f:
                    self.game_stats = json.load(f)
        except:
            pass

    def save_stats(self):
        """保存游戏统计"""
        try:
            with open('game_stats.json', 'w') as f:
                json.dump(self.game_stats, f)
        except:
            pass

    def get_player_color_name(self, player):
        """获取玩家颜色名称"""
        return "Black" if player == 1 else "White"

    def get_player_type_name(self, player):
        """获取玩家类型名称"""
        if player == self.human_player:
            return "Human"
        elif player == self.ai_player:
            return "AI"
        else:
            return f"Player {player}"

    def draw_board(self):
        """绘制棋盘和界面"""
        if not self.gui_enabled:
            return

        self.screen.fill(self.BOARD_COLOR)

        # 绘制网格线
        for i in range(self.board_size):
            # 垂直线
            pygame.draw.line(self.screen, self.LINE_COLOR,
                             (self.padding + i * self.cell_size + self.cell_size // 2,
                              self.padding + self.cell_size // 2),
                             (self.padding + i * self.cell_size + self.cell_size // 2,
                              self.padding + (self.board_size - 1) * self.cell_size + self.cell_size // 2), 2)
            # 水平线
            pygame.draw.line(self.screen, self.LINE_COLOR,
                             (self.padding + self.cell_size // 2,
                              self.padding + i * self.cell_size + self.cell_size // 2),
                             (self.padding + (self.board_size - 1) * self.cell_size + self.cell_size // 2,
                              self.padding + i * self.cell_size + self.cell_size // 2), 2)

        # 绘制天元和星位
        center = self.board_size // 2
        positions = [(center, center)]
        if self.board_size >= 15:
            positions.extend([(3, 3), (3, 11), (11, 3), (11, 11),
                              (3, 7), (7, 3), (7, 11), (11, 7)])

        for row, col in positions:
            if 0 <= row < self.board_size and 0 <= col < self.board_size:
                x = self.padding + col * self.cell_size + self.cell_size // 2
                y = self.padding + row * self.cell_size + self.cell_size // 2
                pygame.draw.circle(self.screen, self.LINE_COLOR, (x, y), 4)

        # 绘制棋子
        last_action = self.last_action

        for r in range(self.board_size):
            for c in range(self.board_size):
                center_x = self.padding + c * self.cell_size + self.cell_size // 2
                center_y = self.padding + r * self.cell_size + self.cell_size // 2
                stone_radius = self.cell_size // 3

                player_on_board = self.env.board[r, c]
                if player_on_board == 1:  # 黑子
                    pygame.draw.circle(self.screen, self.BLACK_STONE_COLOR,
                                       (center_x, center_y), stone_radius)
                    # 高光效果
                    pygame.draw.circle(self.screen, (80, 80, 80),
                                       (center_x - 5, center_y - 5), stone_radius // 4)
                elif player_on_board == -1:  # 白子
                    pygame.draw.circle(self.screen, self.WHITE_STONE_COLOR,
                                       (center_x, center_y), stone_radius)
                    pygame.draw.circle(self.screen, self.LINE_COLOR,
                                       (center_x, center_y), stone_radius, 2)
                    # 高光效果
                    pygame.draw.circle(self.screen, (255, 255, 255),
                                       (center_x - 5, center_y - 5), stone_radius // 4)

                # 标记最后一步
                if last_action is not None and last_action == r * self.board_size + c:
                    pygame.draw.circle(self.screen, self.LAST_MOVE_COLOR,
                                       (center_x, center_y), stone_radius + 3, 3)

        # 绘制状态信息
        self.draw_status_info()
        pygame.display.flip()

    def draw_status_info(self):
        """绘制状态信息"""
        info_y = self.screen_height - 75

        # 游戏状态
        if self.game_over:
            if self.winner == self.human_player:
                message = "🎉 Human Wins!"
                color = (0, 150, 0)
            elif self.winner == self.ai_player:
                message = "🤖 AI Wins!"
                color = (150, 0, 0)
            elif self.winner == 0:
                message = "🤝 Draw!"
                color = (100, 100, 100)
            else:
                message = f"Game Over, Winner: {self.get_player_type_name(self.winner)}"
                color = (0, 0, 0)
            message += " (R to restart)"
        else:
            # 显示当前轮到谁
            player_type = self.get_player_type_name(self.current_turn_player)
            stone_color = self.get_player_color_name(self.current_turn_player)

            if self.current_turn_player == self.human_player:
                color = (0, 100, 0)
                message = f"Your turn ({stone_color})"
            else:
                color = (100, 0, 0)
                message = f"AI's turn ({stone_color})"

        # 主要消息
        text_surface = self.big_font.render(message, True, color)
        text_rect = text_surface.get_rect(center=(self.screen_width // 2, info_y))
        self.screen.blit(text_surface, text_rect)

        # 玩家角色说明
        role_text = f"Human: {self.get_player_color_name(self.human_player)} | " \
                    f"AI: {self.get_player_color_name(self.ai_player)}"
        role_surface = self.font.render(role_text, True, (50, 50, 50))
        role_rect = role_surface.get_rect(center=(self.screen_width // 2, info_y - 25))
        self.screen.blit(role_surface, role_rect)

        # 统计信息
        stats_text = f"Games: {self.game_stats['total_games']} | " \
                     f"Human: {self.game_stats['human_wins']} | " \
                     f"AI: {self.game_stats['ai_wins']} | " \
                     f"Draws: {self.game_stats['draws']}"

        stats_surface = self.font.render(stats_text, True, (100, 100, 100))
        stats_rect = stats_surface.get_rect(center=(self.screen_width // 2, info_y + 25))
        self.screen.blit(stats_surface, stats_rect)

        # 移动计数
        move_text = f"Moves: {self.env.move_count}"
        move_surface = self.font.render(move_text, True, (100, 100, 100))
        move_rect = move_surface.get_rect(center=(self.screen_width // 2, info_y + 45))
        self.screen.blit(move_surface, move_rect)

    def handle_click(self, pos):
        """处理鼠标点击"""
        if self.game_over:
            return None

        x, y = pos
        x -= self.padding
        y -= self.padding

        col = round((x - self.cell_size // 2) / self.cell_size)
        row = round((y - self.cell_size // 2) / self.cell_size)

        if not (0 <= row < self.board_size and 0 <= col < self.board_size):
            return None

        action = row * self.board_size + col
        return action

    def make_move(self, action, player):
        """执行落子"""
        if self.game_over:
            return False

        # 确保轮到当前玩家
        if player != self.current_turn_player:
            print(f"❌ Not {self.get_player_type_name(player)}'s turn!")
            return False

        next_state, reward, done, info = self.env.step(action)

        if info.get('invalid_move', False):
            print(f"❌ Invalid move by {self.get_player_type_name(player)}")
            return False

        self.state = next_state
        self.game_over = done
        self.winner = self.env.winner
        self.last_action = action

        print(f"✅ {self.get_player_type_name(player)} ({self.get_player_color_name(player)}) "
              f"placed stone at ({action // self.board_size}, {action % self.board_size})")

        # 使用环境的当前玩家状态
        if not self.game_over:
            self.current_turn_player = self.env.current_player

        if self.game_over:
            self.update_stats()
            winner_name = self.get_player_type_name(self.winner) if self.winner != 0 else "Draw"
            print(f"🏆 Game Over! Winner: {winner_name}")

        self.draw_board()
        return True

    def update_stats(self):
        """更新游戏统计"""
        self.game_stats['total_games'] += 1

        if self.winner == self.human_player:
            self.game_stats['human_wins'] += 1
        elif self.winner == self.ai_player:
            self.game_stats['ai_wins'] += 1
        elif self.winner == 0:
            self.game_stats['draws'] += 1

        self.save_stats()

    def ai_turn(self):
        """AI回合"""
        if not self.ai_enabled or self.game_over or self.current_turn_player != self.ai_player:
            return

        print(f"🤖 AI ({self.get_player_color_name(self.ai_player)}) is thinking...")

        # 显示思考状态
        if self.show_ai_thinking:
            thinking_surface = self.font.render("AI Thinking...", True, (255, 0, 0))
            thinking_rect = thinking_surface.get_rect(
                center=(self.screen_width // 2, self.padding // 2))
            self.screen.blit(thinking_surface, thinking_rect)
            pygame.display.flip()
            pygame.time.wait(int(self.ai_thinking_time * 1000))

        # AI选择动作
        try:
            ai_action = self.ai_agent.select_action(
                env_state=self.env.get_state(),
                epsilon=0.0  # 测试时不使用随机探索
            )

            print(f"🎯 AI selects action: {ai_action} -> "
                  f"({ai_action // self.board_size}, {ai_action % self.board_size})")
            self.make_move(ai_action, self.ai_player)

        except Exception as e:
            print(f"❌ AI error: {e}")
            legal_actions = self.env.get_legal_actions()
            if legal_actions:
                ai_action = random.choice(legal_actions)
                print(f"🎲 Falling back to random action: {ai_action}")
                self.make_move(ai_action, self.ai_player)
            else:
                print("❌ No legal actions for AI fallback.")

    def reset_game(self):
        """重置游戏"""
        self.state = self.env.reset()
        self.game_over = False
        self.winner = None
        self.last_action = None

        # 重置到初始玩家
        if self.ai_plays_first:
            self.current_turn_player = 1  # AI先手（黑棋）
        else:
            self.current_turn_player = 1  # 人类先手（黑棋）

        print("🔄 Game reset")
        print(f"📋 {self.get_player_type_name(self.current_turn_player)} "
              f"({self.get_player_color_name(self.current_turn_player)}) goes first")
        self.draw_board()

    def play(self):
        """游戏主循环"""
        if not self.gui_enabled:
            print("❌ GUI is not enabled. Cannot start game loop.")
            return

        print("🎮 Starting game loop...")
        clock = pygame.time.Clock()
        running = True
        self.draw_board()

        # 如果AI先手且AI可用
        if self.ai_enabled and self.current_turn_player == self.ai_player:
            pygame.time.wait(500)  # 给玩家时间看到初始状态
            self.ai_turn()

        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False

                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:
                        running = False
                    elif event.key == pygame.K_r:  # R键重置游戏
                        self.reset_game()
                        if self.ai_enabled and self.current_turn_player == self.ai_player:
                            pygame.time.wait(500)
                            self.ai_turn()
                    elif event.key == pygame.K_s and self.ai_enabled and self.ai_agent:  # S键保存模型
                        save_path = self.config.checkpoint_dir
                        if not os.path.exists(save_path):
                            os.makedirs(save_path)
                        model_file_path = os.path.join(save_path,
                                                       f"gui_saved_model_{self.board_size}x{self.board_size}.pth")
                        self.ai_agent.save(model_file_path)
                        print(f"💾 AI Model saved to {model_file_path}")

                elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                    if self.game_over:
                        self.reset_game()
                        if self.ai_enabled and self.current_turn_player == self.ai_player:
                            pygame.time.wait(500)
                            self.ai_turn()
                        continue

                    # 人类玩家回合
                    if self.current_turn_player == self.human_player:
                        action = self.handle_click(event.pos)
                        if action is not None and action in self.env.get_legal_actions():
                            if self.make_move(action, self.human_player):
                                # AI回合
                                if not self.game_over and self.ai_enabled and self.current_turn_player == self.ai_player:
                                    pygame.time.wait(300)  # 短暂延迟
                                    self.ai_turn()
                        else:
                            print("🖱️ Click on a valid empty spot.")

            clock.tick(60)

        print("👋 Exiting game...")
        pygame.quit()
        sys.exit()